adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/__init__.py

@@ -0,0 +1,162 @@
+# ruff: noqa: F403, F401
+from typing import TYPE_CHECKING
+
+from harmony_client import (
+    EvalSample as EvalSample,
+)
+from harmony_client import (
+    EvalSampleInteraction as EvalSampleInteraction,
+)
+from harmony_client import (
+    Grade as Grade,
+)
+from harmony_client import (
+    HarmonyClient as HarmonyClient,
+)
+from harmony_client import (
+    HarmonyJobNotifier as HarmonyJobNotifier,
+)
+from harmony_client import (
+    InferenceModel as InferenceModel,
+)
+from harmony_client import (
+    JobArtifact as JobArtifact,
+)
+from harmony_client import (
+    JobNotifier as JobNotifier,
+)
+from harmony_client import (
+    ModelBuilder as ModelBuilder,
+)
+from harmony_client import (
+    StageNotifier as StageNotifier,
+)
+from harmony_client import (
+    StringThread as StringThread,
+)
+from harmony_client import (
+    TokenizedThread as TokenizedThread,
+)
+from harmony_client import (
+    TrainingModel as TrainingModel,
+)
+from harmony_client import (
+    get_client as get_client,
+)
+from harmony_client import parameters as parameters
+from harmony_client import runtime as runtime
+from rich.progress import Progress
+
+if TYPE_CHECKING:
+    from harmony_client import StringTurn as StringTurn
+else:
+    from typing import NamedTuple
+
+    class StringTurn(NamedTuple):
+        role: str
+        content: str
+
+
+from harmony_client.artifacts.custom_artifact import CustomArtifact
+from harmony_client.artifacts.dataset_artifact import DatasetArtifact
+from harmony_client.file_storage import (
+    FileStorage,
+    FileStorageConfig,
+    LocalFileStorageConfig,
+    S3FileStorageConfig,
+    StoredFile,
+)
+
+import adaptive_harmony.core.rl_utils as rl_utils
+from adaptive_harmony.core.dataset import DataSet
+from adaptive_harmony.core.schedulers import CombinedSchedule, CosineScheduler, CosineSchedulerWithoutWarmup, Scheduler
+from adaptive_harmony.evaluation.evaluation_artifact import EvaluationArtifact
+from adaptive_harmony.metric_logger import Logger, WandbLogger
+
+# Ensure key classes are available at module level
+__all__ = [
+    "StringThread",
+    "StringTurn",
+    "TokenizedThread",
+    "InferenceModel",
+    "ModelBuilder",
+    "TrainingModel",
+    "HarmonyClient",
+    "get_client",
+    "DataSet",
+    "CosineScheduler",
+    "CombinedSchedule",
+    "CosineSchedulerWithoutWarmup",
+    "Scheduler",
+    "WandbLogger",
+    "Logger",
+    "FileStorage",
+    "FileStorageConfig",
+    "LocalFileStorageConfig",
+    "S3FileStorageConfig",
+    "StoredFile",
+    "EvaluationArtifact",
+    "CustomArtifact",
+    "DatasetArtifact",
+    "rl_utils",
+    "Grade",
+    "EvalSample",
+    "EvalSampleInteraction",
+    "JobArtifact",
+]
+
+
+# Patch StringThread to use rich for display
+from harmony_client.runtime.model_artifact_save import save_with_artifact
+
+from adaptive_harmony.core.display import _stringthread_repr, _tokenizedthread_repr
+from adaptive_harmony.core.image_utils import string_thread_to_html_string
+
+# Patch InferenceModel to have json output capabilities
+from adaptive_harmony.core.structured_output import generate_and_validate, render_pydantic_model, render_schema
+
+StringThread.__repr__ = _stringthread_repr  # type: ignore
+TokenizedThread.__repr__ = _tokenizedthread_repr  # type: ignore
+setattr(StringThread, "_repr_html_", string_thread_to_html_string)
+setattr(InferenceModel, "generate_and_validate", generate_and_validate)
+setattr(InferenceModel, "render_schema", staticmethod(render_schema))
+setattr(InferenceModel, "render_pydantic_model", staticmethod(render_pydantic_model))
+
+_original_training_model_save = TrainingModel.save
+
+
+async def _save_with_artifact_wrapper(model: TrainingModel, model_name: str, inference_only: bool = True, ctx=None):
+    return await save_with_artifact(model, model_name, inference_only, ctx, _original_training_model_save)
+
+
+setattr(TrainingModel, "save", _save_with_artifact_wrapper)
+
+
+async def spawn_train(self: ModelBuilder, name: str, max_batch_size: int) -> TrainingModel:
+    fut = await self.spawn_train_with_progress(name, max_batch_size)  # type:ignore
+
+    with Progress() as pbar:
+        task = pbar.add_task("Loading model", total=1000)
+
+        while (prog := await fut._await_progress()) != 1.0:
+            pbar.update(task, completed=prog, total=1.0)
+        pbar.update(task, completed=1.0, total=1.0)
+
+    return await fut.get()
+
+
+async def spawn_inference(self: ModelBuilder, name: str) -> InferenceModel:
+    fut = await self.spawn_inference_with_progress(name)  # type:ignore
+
+    with Progress() as pbar:
+        task = pbar.add_task("Loading model", total=1000)
+
+        while (prog := await fut._await_progress()) != 1.0:
+            pbar.update(task, completed=prog, total=1.0)
+        pbar.update(task, completed=1.0, total=1.0)
+
+    return await fut.get()
+
+
+setattr(ModelBuilder, "spawn_inference", spawn_inference)
+setattr(ModelBuilder, "spawn_train", spawn_train)

adaptive_harmony/common/__init__.py

@@ -0,0 +1,40 @@
+from .callbacks import (
+    CheckpointCallback as CheckpointCallback,
+)
+from .callbacks import (
+    EnvironmentValidationCallback as EnvironmentValidationCallback,
+)
+from .callbacks import (
+    GenerateSamplesCallback as GenerateSamplesCallback,
+)
+from .callbacks import (
+    GraderEvalCallback as GraderEvalCallback,
+)
+from .callbacks import (
+    RecipeCallback as RecipeCallback,
+)
+from .callbacks import (
+    ValidationLossCallback as ValidationLossCallback,
+)
+from .dpo import DPO as DPO
+from .env_grpo import ENVGRPO
+from .grpo import GRPO as GRPO
+from .gspo import GSPO as GSPO
+from .ppo import PPO as PPO
+from .rm import RewardModelling as RewardModelling
+from .sft import SFT as SFT
+
+__all__ = [
+    "SFT",
+    "PPO",
+    "GRPO",
+    "ENVGRPO",
+    "DPO",
+    "RewardModelling",
+    "RecipeCallback",
+    "GenerateSamplesCallback",
+    "ValidationLossCallback",
+    "CheckpointCallback",
+    "GraderEvalCallback",
+    "EnvironmentValidationCallback",
+]

adaptive_harmony/common/callbacks.py

@@ -0,0 +1,219 @@
+from abc import abstractmethod
+from typing import Any
+
+import numpy as np
+from harmony_client import (
+    InferenceModel,
+    StringThread,
+    TrainingModel,
+)
+from loguru import logger
+
+from adaptive_harmony.core.utils import async_map, async_map_fallible
+from adaptive_harmony.environment import EnvironmentFactory
+from adaptive_harmony.graders import BaseGrader
+from adaptive_harmony.logging_table import Table
+
+
+class RecipeCallback:
+    def __init__(self, frequency: float, log_key_prefix: str | None = None):
+        self.frequency = frequency
+        self.last_call = -1.0
+        self.log_key_prefix = log_key_prefix
+
+    async def maybe_call(self, current_percentage: float) -> dict[str, Any]:
+        if current_percentage - self.last_call >= self.frequency:
+            self.last_call = current_percentage
+            callback_dict = await self.callback(current_percentage)
+            prefixed_dict = {
+                (f"{self.log_key_prefix}/{key}" if self.log_key_prefix else key): value
+                for key, value in callback_dict.items()
+            }
+            return prefixed_dict
+        return {}
+
+    @abstractmethod
+    async def callback(self, current_percentage: float) -> dict[str, Any]: ...
+
+
+class GenerateSamplesCallback(RecipeCallback):
+    def __init__(
+        self,
+        thread_set: list[StringThread],
+        model: InferenceModel,
+        frequency: float,
+        log_key: str = "samples",
+    ):
+        super().__init__(frequency, log_key_prefix="generation")
+        self.thread_set = thread_set
+        self.model = model
+        self.log_key = log_key
+
+    async def callback(self, current_percentage: float) -> dict[str, Any]:
+        logger.info("Entering generation callback...")
+        generation_tokens = await async_map_fallible(self.model.generate_tokens, self.thread_set)
+        generation_results = await async_map_fallible(self.model.detokenize_thread, generation_tokens)
+        gen_lengths = [sample.len_last_turn() for sample in generation_tokens]
+
+        generation_logs = {
+            self.log_key: Table()
+            .add_column(
+                "system",
+                [
+                    sample.get_turns()[0].content if sample.get_turns()[0].role == "system" else ""
+                    for sample in generation_results
+                ],
+            )
+            .add_column(
+                "prompt",
+                [
+                    repr(
+                        StringThread(
+                            sample.get_turns()[1:-1]
+                            if (sample.get_turns() and sample.get_turns()[0].role == "system")
+                            else sample.get_turns()[:-1]
+                        )
+                    )
+                    for sample in generation_results
+                ],
+            )
+            .add_column("response", [response.last_content() for response in generation_results]),
+            "generation_length_mean": np.mean(gen_lengths).item(),
+            "generation_length_std": np.std(gen_lengths).item(),
+            "num_samples": len(generation_results),
+        }
+        return generation_logs
+
+
+class ValidationLossCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        frequency: float = 0.1,
+        log_key: str = "loss",
+    ):
+        super().__init__(frequency, log_key_prefix="validation")
+        self.validation_set = validation_set
+        self.model = model
+        self.log_key = log_key
+
+    async def callback(self, current_percentage: float) -> dict[str, float]:
+        logger.info("Entering validation loss callback...")
+        losses = []
+        tokens = await async_map_fallible(self.model.tokenize_thread, self.validation_set)
+        logprobs = await async_map(self.model.logprobs_per_token, tokens)
+        losses = [-(sum(lp) / len(lp)) for lp in logprobs]
+
+        return {self.log_key: sum(losses) / len(losses)}
+
+
+class CheckpointCallback(RecipeCallback):
+    def __init__(
+        self,
+        model: TrainingModel,
+        checkpoint_name: str,
+        frequency: float = 0.2,
+    ):
+        super().__init__(frequency, log_key_prefix="checkpointing")
+        self.last_call = 0.0  # avoid saving the model at the first period
+        self.model = model
+        self.model_log_name = checkpoint_name
+
+    async def callback(self, current_percentage: float):
+        logger.info(f"Saving checkpoint at {current_percentage * 100} % of training ...")
+        await self.model.save(f"{self.model_log_name}-{round(current_percentage, 3)}")
+        return {}
+
+
+class GraderEvalCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        grader: BaseGrader,
+        frequency: float,
+        log_key: str = "validation",
+        clear_grader_logs: bool = True,
+        temperature: float = 0.0,
+    ):
+        super().__init__(frequency, log_key_prefix=log_key)
+        self.validation_set = validation_set
+        self.model = model
+        self.grader = grader
+        self.clear_grader_logs = clear_grader_logs
+        self.temperature = temperature
+
+    async def callback(self, current_percentage: float) -> dict[str, float | Table]:
+        logger.info("Entering grader evaluation callback...")
+        temp_model = self.model.temperature(self.temperature)
+
+        tokenized_results = await async_map_fallible(temp_model.generate_tokens, self.validation_set)
+        string_results = await async_map(temp_model.detokenize_thread, tokenized_results)
+        grades = await async_map_fallible(self.grader.grade, string_results)
+        gen_lengths = [sample.len_last_turn() for sample in tokenized_results]
+
+        grader_logs = self.grader.get_logs(clear=self.clear_grader_logs)
+        return {
+            **{f"rewards/{key}": value for key, value in grader_logs.items()},
+            "generation_length_mean": float(np.mean(gen_lengths).item()),
+            "generation_length_std": float(np.std(gen_lengths).item()),
+            "num_samples": float(len(grades)),
+        }
+
+
+class EnvironmentValidationCallback(RecipeCallback):
+    def __init__(
+        self,
+        validation_set: list[StringThread],
+        model: InferenceModel,
+        env_factory: EnvironmentFactory,
+        frequency: float,
+        log_key: str = "validation",
+        clear_env_logs: bool = True,
+        temperature: float = 0.0,
+        num_samples_log: int = 0,
+    ):
+        super().__init__(frequency, log_key_prefix=log_key)
+        self.validation_set = validation_set
+        self.model = model
+        self.env_factory = env_factory
+        self.clear_env_logs = clear_env_logs
+        self.temperature = temperature
+        self.num_samples_log = num_samples_log
+
+    async def generate_trajectory(self, initial_thread: StringThread) -> tuple[StringThread, float, int]:
+        env = self.env_factory.create_environment(initial_thread.metadata)
+        temp_model = self.model.temperature(self.temperature)
+        trajectory, trajectory_score = await env.generate_trajectory_and_grade(temp_model, initial_thread)
+        num_turns = len([turn for turn in trajectory.get_turns() if turn.role == "assistant"])
+        return trajectory, trajectory_score.cumulative_score, num_turns
+
+    async def callback(self, current_percentage: float) -> dict[str, float | Table]:
+        logger.info("Entering environment validation callback...")
+
+        results = await async_map_fallible(self.generate_trajectory, self.validation_set)
+
+        trajectories = [traj for traj, _, _ in results]
+        scores = [score for _, score, _ in results]
+        num_turns_list = [num_turns for _, _, num_turns in results]
+
+        validation_logs = {
+            "score_mean": np.mean(scores).item(),
+            "score_std": np.std(scores).item(),
+            "num_turns_mean": np.mean(num_turns_list).item(),
+            "num_turns_std": np.std(num_turns_list).item(),
+            "num_samples": len(results),
+        }
+
+        env_logs = self.env_factory.get_logs(clear=self.clear_env_logs)
+        validation_logs.update({f"env/{key}": value for key, value in env_logs.items()})
+
+        if self.num_samples_log > 0:
+            samples = [repr(traj) for traj in trajectories[: self.num_samples_log]]
+            samples_scores = scores[: self.num_samples_log]
+            table = Table().add_column("trajectory", samples).add_column("score", samples_scores)
+            validation_logs["samples"] = table
+
+        logger.info(f"Validation Mean score: {validation_logs['score_mean']:.4f}")
+        return validation_logs

adaptive_harmony/common/checkpointing.py

@@ -0,0 +1,163 @@
+import json
+import os
+from pathlib import Path
+from typing import Awaitable, Callable, Sequence
+
+import anyio
+import numpy as np
+from loguru import logger as loguru
+
+from adaptive_harmony import DataSet, StringThread
+from adaptive_harmony.common.callbacks import RecipeCallback
+from adaptive_harmony.core.utils import hash_dataset
+
+
+class CheckpointManager:
+    def __init__(
+        self,
+        recipe_name: str,
+        dataset: DataSet,
+        threads_dataset: Sequence[StringThread],
+        callbacks: Sequence[RecipeCallback],
+        hyperparams_hash: str,
+        job_id: str | None = None,
+        checkpoint_frequency: float = 0.2,
+        restart_from_checkpoint: str | None = None,
+    ):
+        self.recipe_name = recipe_name
+        self.dataset = dataset
+        self.dataset_hash = hash_dataset(threads_dataset)
+        self.hyperparams_hash = hyperparams_hash
+        self.callbacks = callbacks
+        self.checkpoint_frequency = checkpoint_frequency
+        self.last_checkpoint_percentage = 0.0
+        self.restart_from_checkpoint = restart_from_checkpoint
+        self.job_id = job_id
+        self.checkpointing_folder = self._init_folder()
+
+    def _init_folder(self) -> str | None:
+        if self.job_id is None or os.getenv("HARMONY_NO_CHECKPOINTING") is not None:
+            loguru.warning("Checkpointing is disabled for this recipe.")
+            return None
+        return os.path.join(os.getenv("RECIPE_CHECKPOINTS_DIR", "/checkpoints"), self.job_id)
+
+    async def maybe_restore_checkpoint(
+        self,
+        recipe_specific_checkpoint_loading: Callable[[dict], Awaitable[None]],
+    ) -> None:
+        if self.restart_from_checkpoint is None:
+            return
+
+        checkpoint_path = Path(self.restart_from_checkpoint)
+        checkpoint_file = self._resolve_checkpoint_file(checkpoint_path)
+
+        assert checkpoint_file, f"Checkpoint file not found: {checkpoint_path}."
+
+        loguru.info(f"Loading {self.recipe_name} checkpoint from: {checkpoint_file}")
+
+        contents = ""
+        async with await anyio.open_file(checkpoint_file, "r") as f:
+            contents = await f.read()
+        checkpoint_data = json.loads(contents)
+
+        assert checkpoint_data.get("recipe_type") == self.recipe_name, (
+            f"Recipe type mismatch: checkpoint is '{checkpoint_data.get('recipe_type')}', "
+            f"but trying to load into {self.recipe_name}"
+        )
+
+        assert checkpoint_data.get("dataset_hash") == self.dataset_hash, (
+            "Dataset hash mismatch between checkpoint and current dataset."
+        )
+
+        assert checkpoint_data.get("hyperparams_hash") == self.hyperparams_hash, (
+            "Hyperparameters hash mismatch between checkpoint and current recipe configuration."
+        )
+
+        self.dataset.idx = checkpoint_data.get("dataset_idx", 0)
+
+        access_indices_list = checkpoint_data.get("dataset_access_indices", [])
+        if access_indices_list:
+            self.dataset.access_indices = np.array(access_indices_list)
+
+        rng_state = checkpoint_data.get("dataset_rng_state")
+        if rng_state:
+            self.dataset.rng.bit_generator.state = rng_state
+
+        callback_states = checkpoint_data.get("callback_last_calls", [])
+        assert len(callback_states) == len(self.callbacks), "Mismatch in number of callbacks when loading checkpoint"
+        for i, callback in enumerate(self.callbacks):
+            callback.last_call = callback_states[i]
+
+        await recipe_specific_checkpoint_loading(checkpoint_data)
+
+        self.last_checkpoint_percentage = checkpoint_data.get("completion_percentage", 0.0)
+
+        loguru.info(f"Checkpoint restored: starting {self.recipe_name} from {self.last_checkpoint_percentage:.2%}.")
+
+    async def maybe_checkpoint(
+        self,
+        completion_percentage: float,
+        recipe_specific_checkpoint_saving: Callable[[], Awaitable[dict]],
+    ) -> bool:
+        if self.checkpointing_folder is None:
+            return False
+
+        if completion_percentage >= 1.0:
+            return False
+
+        if await self._check_graceful_exit_file():
+            loguru.info(f"Graceful exit requested. Saving checkpoint and exiting {self.recipe_name} training loop.")
+            await self._save_checkpoint(completion_percentage, recipe_specific_checkpoint_saving)
+            return True
+
+        if completion_percentage - self.last_checkpoint_percentage >= self.checkpoint_frequency:
+            await self._save_checkpoint(completion_percentage, recipe_specific_checkpoint_saving)
+            self.last_checkpoint_percentage = completion_percentage
+
+        return False
+
+    async def _save_checkpoint(
+        self,
+        completion_percentage: float,
+        get_save_config: Callable[[], Awaitable[dict]],
+    ) -> None:
+        assert self.checkpointing_folder is not None  # will never be called outside of this condition
+        progress_pct = int(completion_percentage * 100)
+        checkpoint_dir = Path(self.checkpointing_folder)
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+        loguru.info(f"Checkpointing {self.recipe_name} at {checkpoint_dir} ({progress_pct}%)...")
+
+        recipe_data = await get_save_config()
+
+        checkpoint_data = {
+            "recipe_type": self.recipe_name,
+            "dataset_hash": self.dataset_hash,
+            "hyperparams_hash": self.hyperparams_hash,
+            "dataset_idx": self.dataset.idx,
+            "dataset_access_indices": self.dataset.access_indices.tolist(),
+            "dataset_rng_state": self.dataset.rng.bit_generator.state,
+            "callback_last_calls": [callback.last_call for callback in self.callbacks],
+            "completion_percentage": completion_percentage,
+            **recipe_data,
+        }
+
+        checkpoint_file = checkpoint_dir / f"checkpoint-{progress_pct}.json"
+
+        data_dump = json.dumps(checkpoint_data, indent=2)
+        async with await anyio.open_file(checkpoint_file, "w") as f:
+            await f.write(data_dump)
+
+        loguru.info(f"Checkpoint saved: {checkpoint_file}")
+
+    async def _check_graceful_exit_file(self) -> bool:
+        if self.checkpointing_folder is None:
+            return False
+        return (Path(self.checkpointing_folder) / "GRACEFUL_EXIT").exists()
+
+    @staticmethod
+    def _resolve_checkpoint_file(path: Path) -> Path | None:
+        if path.is_dir():
+            files = sorted(path.glob("checkpoint-*.json"), key=lambda p: int(p.stem.split("-")[1]))
+            return files[-1] if files else None
+        return path if path.exists() else None

adaptive_harmony/common/dpo.py

@@ -0,0 +1,92 @@
+from typing import Callable, Sequence
+
+from tqdm.auto import tqdm
+
+from adaptive_harmony import (
+    CosineScheduler,
+    DataSet,
+    JobNotifier,
+    Logger,
+    StageNotifier,
+    StringThread,
+    TrainingModel,
+)
+from adaptive_harmony.common.callbacks import RecipeCallback
+from adaptive_harmony.core.utils import async_map_batch, log_args
+from adaptive_harmony.metric_logger import StdoutLogger
+
+
+class DPO:
+    @log_args
+    def __init__(
+        self,
+        dataset: list[tuple[StringThread, StringThread]],  # (positive_sample, negative_sample)
+        model: TrainingModel,
+        logger: Logger = StdoutLogger(),
+        stage_notifier: StageNotifier = JobNotifier().stage_notifier("DPO Training"),
+        callbacks: Sequence[RecipeCallback] = [],
+        lr: float = 1e-6,
+        lr_scheduler: Callable[[float], float] | None = None,
+        samples_per_batch=32,
+        max_grad_norm=1.0,
+        kl_beta=0.1,
+        epochs=1,
+        skip_nan_gradients: bool = False,
+    ):
+        # Core components
+        self.model_ref = None
+        self.dataset = DataSet(dataset)
+        self.model = model
+        self.logger = logger
+        self.stage_notifier = stage_notifier
+        self.callbacks = callbacks
+        self.lr_schedule = lr_scheduler or CosineScheduler(lr)
+        self.samples_per_batch = samples_per_batch
+        self.max_grad_norm = max_grad_norm
+        self.skip_nan_gradients = skip_nan_gradients
+
+        # DPO HP's
+        self.kl_beta = kl_beta
+        self.epochs = epochs
+
+    @property
+    def training_completion_percentage(self):
+        return self.dataset.completion_percentage() / self.epochs
+
+    async def process_sample(self, sample: tuple[StringThread, StringThread]):
+        assert self.model_ref is not None, "Calling `process_sample_dpo` before reference model has been set"
+
+        pos, neg = sample
+        ref_logprobs_pos = await self.model_ref.logprobs(pos)
+        ref_logprobs_neg = await self.model_ref.logprobs(neg)
+        await self.model.train_dpo(pos, neg, ref_logprobs_pos, ref_logprobs_neg, self.kl_beta)
+
+    async def run(self):
+        self.model_ref = await self.model.clone_inf()
+
+        self.stage_notifier.report_progress(
+            tot_num_samples=len(self.dataset) * self.epochs,
+            processed_num_samples=self.dataset.idx,
+            monitoring_link=self.logger.training_monitoring_link,
+        )
+
+        with tqdm(total=100) as pbar:
+            while self.training_completion_percentage < 1.0:
+                for callback in self.callbacks:
+                    if logs := await callback.maybe_call(self.training_completion_percentage):
+                        self.logger(logs)
+
+                await async_map_batch(self.process_sample, self.dataset, self.samples_per_batch)
+                cp = self.training_completion_percentage
+                current_lr = self.lr_schedule(cp)
+                pbar.update(cp * 100.0 - pbar.n)
+                logs = await self.model.optim_step(
+                    current_lr, wd=0, max_grad_norm=self.max_grad_norm, skip_nan_gradients=self.skip_nan_gradients
+                )
+                self.logger(logs | dict(completion_percentage=cp))
+
+                self.stage_notifier.report_progress(
+                    tot_num_samples=len(self.dataset) * self.epochs,
+                    processed_num_samples=self.dataset.idx,
+                    monitoring_link=self.logger.training_monitoring_link,
+                )